'''
Function:
    Build the segmentor
Author:
    Zhenchao Jin
'''
import copy
from .encnet import ENCNet
from .upernet import UPerNet
from .deeplabv3plus import Deeplabv3Plus
from .fcn import FCN, DepthwiseSeparableFCN


'''BuildSegmentor'''
def BuildSegmentor(segmentor_cfg, losses_cfg, mode):
    supported_segmentors = {
        'fcn': FCN,
        'encnet': ENCNet,
        'upernet': UPerNet,
        'deeplabv3plus': Deeplabv3Plus,
    }
    selected_segmentor = supported_segmentors[segmentor_cfg['type']]
    segmentor_cfg = copy.deepcopy(segmentor_cfg)
    segmentor_cfg.pop('type')
    return selected_segmentor(segmentor_cfg, losses_cfg, mode)
    # return selected_segmentor(segmentor_cfg, losses_cfg, mode='TRAIN')